#include "njm_device.hpp"
#include <stdexcept>
#include <vector>
#include <iostream>
#include <cassert>

namespace njm
{

    NjmDevice::NjmDevice(NjmWindow &window)
    {
        VkApplicationInfo appInfo{};
        appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
        appInfo.apiVersion = VK_API_VERSION_1_4;
        appInfo.applicationVersion = VK_MAKE_API_VERSION(0, 1, 0, 0);
        appInfo.pApplicationName = window.getWindowName().c_str();
        appInfo.engineVersion = VK_MAKE_API_VERSION(0, 0, 0, 0);
        appInfo.pEngineName = "NO ENGINE";

        mInstance = createInstance(appInfo);
        window.createWindowSurface(mInstance, &mSurface);

        auto devTuple = selectPhysicalDevice(mInstance);
        mPhysicalDevice = std::get<0>(devTuple);
        auto queueCreateInfo = std::get<1>(devTuple);
        myVkDevice = createLogicalDevice(mPhysicalDevice, queueCreateInfo);
    }

    NjmDevice::~NjmDevice()
    {
        vkDestroyCommandPool(myVkDevice, getCommandPool(), nullptr);
        if(mSurface) vkDestroySurfaceKHR(mInstance, mSurface, nullptr);
        if(myVkDevice) vkDestroyDevice(myVkDevice, nullptr);
        if(mInstance) vkDestroyInstance(mInstance, nullptr);
    }

    VkInstance NjmDevice::createInstance(VkApplicationInfo appInfo)
    {
        uint32_t glfwExCnt = 0;
        const char **glfwExs = glfwGetRequiredInstanceExtensions(&glfwExCnt);
        for (size_t i = 0; i < glfwExCnt; i++)
        {
            std::cout << "glfw required extensions: " << glfwExs[i] << std::endl;
        }
        // const std::vector<const char*> desiredExtensions = {
        //     "VK_KHR_surface",
        //     "VK_KHR_win32_surface", "VK_KHR_xlib_surface", "VK_KHR_wayland_surface", "VK_KHR_android_surface",
        //     "VK_EXT_debug_utils",
        //     "VK_KHR_timeline_semaphore",
        //     "VK_KHR_synchronization2",
        //     "VK_KHR_get_physical_device_properties2",
        //     "VK_KHR_external_memory_capabilities",
        //     "VK_KHR_external_semaphore_capabilities",
        // };
        const std::vector<const char *> desiredInstanceExtensions = {"VK_KHR_surface", "VK_KHR_win32_surface", "VK_EXT_debug_utils"};
        const auto exCnt = desiredInstanceExtensions.size();
        checkInstanceExtensions(desiredInstanceExtensions);

        // const std::vector<const char*> desiredLayers = {
        //     "VK_LAYER_KHRONOS_validation",
        //     "VK_LAYER_LUNARG_gfxreconstruct",
        //     "VK_LAYER_LUNARG_api_dump",
        //     "VK_LAYER_KHRONOS_profiles",
        //     "VK_LAYER_LUNARG_monitor",
        //     "VK_LAYER_LUNARG_screenshot",
        // };
        const std::vector<const char *> desiredLayers = {"VK_LAYER_KHRONOS_validation"};
        const auto layerCnt = desiredLayers.size();
        checkInstanceLayers(desiredLayers);

        VkInstanceCreateInfo insInfo{};
        insInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
        insInfo.pApplicationInfo = &appInfo;
        insInfo.enabledExtensionCount = static_cast<uint32_t>(exCnt);
        insInfo.ppEnabledExtensionNames = desiredInstanceExtensions.data();
        insInfo.enabledLayerCount = static_cast<uint32_t>(layerCnt);
        insInfo.ppEnabledLayerNames = desiredLayers.data();

        VkInstance ins;
        if (vkCreateInstance(&insInfo, nullptr, &ins) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to create vk instance");
        }
        return ins;
    }

    std::tuple<VkPhysicalDevice, VkDeviceQueueCreateInfo> NjmDevice::selectPhysicalDevice(VkInstance instance)
    {
        uint32_t devCnt = 0;
        vkEnumeratePhysicalDevices(instance, &devCnt, nullptr);
        if (devCnt == 0)
        {
            throw std::runtime_error("failed to find gpu with vk-support");
        }
        std::vector<VkPhysicalDevice> devs(devCnt);
        vkEnumeratePhysicalDevices(instance, &devCnt, devs.data());

        for (auto &&dev : devs)
        {
            uint32_t propCnt = 0;
            vkGetPhysicalDeviceQueueFamilyProperties(dev, &propCnt, nullptr);
            if (propCnt == 0)
            {
                continue;
            }
            std::vector<VkQueueFamilyProperties> props(propCnt);
            vkGetPhysicalDeviceQueueFamilyProperties(dev, &propCnt, props.data());
            VkPhysicalDeviceProperties devProps{};
            vkGetPhysicalDeviceProperties(dev, &devProps);

            for (uint32_t i = 0; i < propCnt; i++)
            {
                if (props[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
                {

                    mGraphicsQueueFamilyIndex = i;
                    mPhysicalDeviceProperties = devProps;
                    mQueueFamilyProperties = props[i];

                    float queuePriority = 1.0f;
                    VkDeviceQueueCreateInfo queueCreateInfo{};
                    queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
                    queueCreateInfo.queueFamilyIndex = mGraphicsQueueFamilyIndex;
                    queueCreateInfo.queueCount = props[i].queueCount;
                    queueCreateInfo.pQueuePriorities = &queuePriority;
                    return std::make_tuple(dev, queueCreateInfo);
                }
            }
        }
        throw std::runtime_error("no suitable gpu found");
    }

    VkDevice NjmDevice::createLogicalDevice(VkPhysicalDevice physicalDevice, VkDeviceQueueCreateInfo queueCreateInfo)
    {
        // VK_KHR_SWAPCHAIN_EXTENSION_NAME : VK_KHR_swapchain
        const std::vector<const char *> desiredDeviceExtensions = {VK_KHR_SWAPCHAIN_EXTENSION_NAME};
        checkDeviceExtensions(physicalDevice, desiredDeviceExtensions);
        auto exCnt = desiredDeviceExtensions.size();

        VkPhysicalDeviceFeatures deviceFeatures = {};
        vkGetPhysicalDeviceFeatures(physicalDevice, &deviceFeatures);

        VkDeviceCreateInfo deviceCreateInfo{};
        deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        deviceCreateInfo.queueCreateInfoCount = 1,
        deviceCreateInfo.pQueueCreateInfos = &queueCreateInfo;
        deviceCreateInfo.enabledExtensionCount = static_cast<uint32_t>(exCnt);
        deviceCreateInfo.ppEnabledExtensionNames = desiredDeviceExtensions.data();
        // deviceCreateInfo.enabledLayerCount = ;
        // deviceCreateInfo.ppEnabledLayerNames = ;
        // deviceCreateInfo.pEnabledFeatures = ;

        VkDevice device;
        if (vkCreateDevice(physicalDevice, &deviceCreateInfo, nullptr, &device) != VK_SUCCESS)
        {
            throw std::runtime_error("Failed to create logical device!");
        }
        return device;
    }

    //
    VkDevice NjmDevice::device() { return myVkDevice; }

    VkCommandPool NjmDevice::getCommandPool()
    {
        static VkCommandPool mCommandPool; // 静态变量确保单例模式
        if (mCommandPool == nullptr)
        {
            VkCommandPoolCreateInfo createInfo = {};
            createInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
            createInfo.queueFamilyIndex = mGraphicsQueueFamilyIndex;            // 假设device构建时，已通过findQueueFamily获取
            createInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; // 允许重置命令缓冲区

            if (vkCreateCommandPool(myVkDevice, &createInfo, nullptr, &mCommandPool) != VK_SUCCESS)
            {
                throw std::runtime_error("Failed to create command pool!");
            }
        }
        return mCommandPool;
    }

    uint32_t NjmDevice::findMemoryType(uint32_t requirementsMemoryTypeBits, VkMemoryPropertyFlags properties)
    {
        VkPhysicalDeviceMemoryProperties memProperties;
        vkGetPhysicalDeviceMemoryProperties(mPhysicalDevice, &memProperties);
        for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++)
        {
            // 检查内存类型位是否匹配，并且属性包含所需属性
            if ((requirementsMemoryTypeBits & (1 << i)) &&
                (memProperties.memoryTypes[i].propertyFlags & properties) == properties)
            {
                return i;
            }
        }
        assert(0 && "Failed to find suitable memory type!"); // 如果未找到合适的类型，按规范应至少存在一个可用类型，此处处理错误
        return 0;                                            // 默认返回0（实际应避免到达此处）
    }

    uint32_t NjmDevice::findQueueFamily(VkPhysicalDevice physicalDevice, VkSurfaceKHR surface)
    {
        uint32_t queueFamilyCount = 0;
        vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, nullptr);
        std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
        vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, queueFamilies.data());

        int32_t graphicsFamily = -1;
        int32_t presentFamily = -1;
        for (uint32_t i = 0; i < queueFamilyCount; ++i)
        {
            // 检查是否支持图形操作
            if (queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
            {
                graphicsFamily = i;
            }
            // 检查是否支持呈现到表面
            VkBool32 presentSupport = false;
            vkGetPhysicalDeviceSurfaceSupportKHR(physicalDevice, i, surface, &presentSupport);
            if (presentSupport)
            {
                presentFamily = i;
            }
            // 如果找到同时支持图形和呈现的队列族，直接返回
            if (graphicsFamily != -1 && presentFamily != -1)
            {
                return graphicsFamily; // 返回同时支持两者的队列族索引
            }
        }

        // 如果未找到同时支持的队列族，尝试分开处理
        if (graphicsFamily == -1)
        {
            throw std::runtime_error("Failed to find a graphics queue family!");
        }
        if (presentFamily == -1)
        {
            throw std::runtime_error("Failed to find a present queue family!");
        }
        // 如果队列族分离，可能需要合并使用（但通常建议使用同一队列族以避免数据传输开销）
        // 此处抛出错误，因为大多数情况下应存在同时支持的队列族
        throw std::runtime_error("Failed to find a queue family that supports both graphics and presentation!");
    }

    void NjmDevice::createBuffer(VkDeviceSize size, VkBufferUsageFlags usage,
                                 VkMemoryPropertyFlags properties, VkBuffer &buffer, VkDeviceMemory &bufferMemory)
    {
        VkBufferCreateInfo bufferInfo{};
        bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
        bufferInfo.size = size;
        bufferInfo.usage = usage;
        bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
        if (vkCreateBuffer(myVkDevice, &bufferInfo, nullptr, &buffer) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to create buffer");
        }

        VkMemoryRequirements requirements;
        vkGetBufferMemoryRequirements(myVkDevice, buffer, &requirements);

        VkMemoryAllocateInfo allocInfo{};
        allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
        allocInfo.allocationSize = requirements.size;
        allocInfo.memoryTypeIndex = findMemoryType(requirements.memoryTypeBits, properties);
        if (vkAllocateMemory(myVkDevice, &allocInfo, nullptr, &bufferMemory) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to allocate buffer memory");
        }
        vkBindBufferMemory(myVkDevice, buffer, bufferMemory, 0);
    }

    void NjmDevice::createImageWithInfo(VkImageCreateInfo imgInfo, VkMemoryPropertyFlagBits memoryPropertyFlag, VkImage &image, VkDeviceMemory &imgMemory)
    {
        if (vkCreateImage(myVkDevice, &imgInfo, nullptr, &image) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to create image");
        }

        VkMemoryRequirements reqs;
        vkGetImageMemoryRequirements(myVkDevice, image, &reqs);
        VkMemoryAllocateInfo allocInfo{};
        allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
        allocInfo.allocationSize = reqs.size;
        allocInfo.memoryTypeIndex = findMemoryType(reqs.memoryTypeBits, memoryPropertyFlag);
        if (vkAllocateMemory(myVkDevice, &allocInfo, nullptr, &imgMemory) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to allocate image memory");
        }
        vkBindImageMemory(myVkDevice, image, imgMemory, 0);
    }

    void NjmDevice::copyBuffer(VkBuffer srcBuf, VkBuffer dstBuf, VkDeviceSize bufSize)
    {
        assert(srcBuf != VK_NULL_HANDLE && dstBuf != VK_NULL_HANDLE);
        // 创建一次性命令缓冲区
        VkCommandBufferAllocateInfo allocInfo = {};
        allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        allocInfo.commandPool = getCommandPool();
        allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        allocInfo.commandBufferCount = 1;
        VkCommandBuffer cmdBuffer;
        vkAllocateCommandBuffers(myVkDevice, &allocInfo, &cmdBuffer);

        // 记录命令缓冲区
        VkCommandBufferBeginInfo beginInfo = {};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
        vkBeginCommandBuffer(cmdBuffer, &beginInfo);

        // 添加缓冲区复制命令
        VkBufferCopy copyRegion = {};
        copyRegion.srcOffset = 0;  // 源缓冲区偏移
        copyRegion.dstOffset = 0;  // 目标缓冲区偏移
        copyRegion.size = bufSize; // 复制大小
        vkCmdCopyBuffer(cmdBuffer, srcBuf, dstBuf, 1, &copyRegion);

        // 添加内存屏障（可选）, 如果需要内存屏障（例如后续立即使用目标缓冲区）
        // VkMemoryBarrier barrier = {};
        // barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
        // barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        // barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        // vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_VERTEX_INPUT_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);

        vkEndCommandBuffer(cmdBuffer);

        // 提交命令缓冲区
        VkSubmitInfo submitInfo = {};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers = &cmdBuffer;

        VkQueue q{};
        vkGetDeviceQueue(myVkDevice, mGraphicsQueueFamilyIndex, 0, &q);
        vkQueueSubmit(q, 1, &submitInfo, VK_NULL_HANDLE);
        vkQueueWaitIdle(q); // 等待队列完成
        // 释放命令缓冲区
        vkFreeCommandBuffers(myVkDevice, getCommandPool(), 1, &cmdBuffer);
    }

    void NjmDevice::copyBufferToImage(VkBuffer buffer, VkImage image, uint32_t width, uint32_t height)
    {
        VkCommandBufferAllocateInfo allocInfo = {};
        allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        allocInfo.commandPool = getCommandPool();
        allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        allocInfo.commandBufferCount = 1;
        VkCommandBuffer cmdBuffer;
        vkAllocateCommandBuffers(myVkDevice, &allocInfo, &cmdBuffer);
        VkCommandBufferBeginInfo beginInfo = {};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
        vkBeginCommandBuffer(cmdBuffer, &beginInfo);

        VkImageSubresourceRange range = {};
        range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        range.baseMipLevel = 0;
        range.levelCount = 1;
        range.baseArrayLayer = 0;
        range.layerCount = 1;
        VkImageMemoryBarrier barrier1 = {};
        barrier1.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        barrier1.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        barrier1.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        barrier1.srcAccessMask = 0;
        barrier1.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        barrier1.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier1.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier1.image = image;
        barrier1.subresourceRange = range;
        vkCmdPipelineBarrier(cmdBuffer, 
            VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 
            0, 0, nullptr, 0, nullptr, 1, &barrier1);

        VkBufferImageCopy region = {};
        region.bufferOffset = 0;
        region.bufferRowLength = 0;
        region.bufferImageHeight = 0;
        region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        region.imageSubresource.mipLevel = 0;
        region.imageSubresource.baseArrayLayer = 0;
        region.imageSubresource.layerCount = 1;
        region.imageOffset = {0, 0, 0};
        region.imageExtent = {width, height, 1};
        vkCmdCopyBufferToImage(cmdBuffer, buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);
        
        VkImageMemoryBarrier barrier2 = {};
        barrier2.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        barrier2.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        barrier2.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        barrier2.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier2.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier2.image = image;
        barrier2.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        barrier2.subresourceRange.baseMipLevel = 0;
        barrier2.subresourceRange.levelCount = 1;
        barrier2.subresourceRange.baseArrayLayer = 0;
        barrier2.subresourceRange.layerCount = 1;
        barrier2.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        barrier2.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        vkCmdPipelineBarrier(cmdBuffer, 
            VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 
            0, 0, nullptr, 0, nullptr, 1, &barrier2);

        vkEndCommandBuffer(cmdBuffer);
        VkSubmitInfo submitInfo = {};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers = &cmdBuffer;
        VkQueue q{};
        vkGetDeviceQueue(myVkDevice, mGraphicsQueueFamilyIndex, 0, &q);
        vkQueueSubmit(q, 1, &submitInfo, VK_NULL_HANDLE);
        vkQueueWaitIdle(q);
        vkFreeCommandBuffers(myVkDevice, getCommandPool(), 1, &cmdBuffer);
    }

    // helpers
    void NjmDevice::checkInstanceExtensions(const std::vector<const char *> &names)
    {
        uint32_t count;
        vkEnumerateInstanceExtensionProperties(nullptr, &count, nullptr);
        std::vector<VkExtensionProperties> props(count);
        vkEnumerateInstanceExtensionProperties(nullptr, &count, props.data());
        std::cout << "supported instance-extensions: " << std::endl;
        // for (auto &&supportedEx : props)
        // {
        //     std::cout << supportedEx.extensionName << ", ";
        // }
        // std::cout << std::endl;
        for (auto &name : names)
        {
            bool found = false;
            for (auto &prop : props)
            {
                if (strcmp(name, prop.extensionName) == 0)
                {
                    found = true;
                    break;
                }
            }
            if (!found)
            {
                throw std::runtime_error("desired vkInstance-extension not found: " + std::string(name));
            }
        }
    }

    void NjmDevice::checkInstanceLayers(const std::vector<const char *> &names)
    {
        uint32_t count;
        vkEnumerateInstanceLayerProperties(&count, nullptr);
        std::vector<VkLayerProperties> props(count);
        vkEnumerateInstanceLayerProperties(&count, props.data());
        std::cout << "supported layers: " << std::endl;
        // for (auto &&supportedLayer : props)
        // {
        //     std::cout << supportedLayer.layerName << ", ";
        // }
        // std::cout << std::endl;
        for (auto &name : names)
        {
            bool found = false;
            for (auto &prop : props)
            {
                if (strcmp(name, prop.layerName) == 0)
                {
                    found = true;
                    break;
                }
            }
            if (!found)
            {
                throw std::runtime_error("desired vkInstance-layer not found: " + std::string(name));
            }
        }
    }

    void NjmDevice::checkDeviceExtensions(VkPhysicalDevice physicalDevice, const std::vector<const char *> &names)
    {
        uint32_t count;
        vkEnumerateDeviceExtensionProperties(physicalDevice, nullptr, &count, nullptr);
        std::vector<VkExtensionProperties> props(count);
        vkEnumerateDeviceExtensionProperties(physicalDevice, nullptr, &count, props.data());
        std::cout << "supported device-extensions: " << std::endl;
        // for (auto &&supportedEx : props)
        // {
        //     std::cout << supportedEx.extensionName << ", ";
        // }
        // std::cout << std::endl;
        for (auto &name : names)
        {
            bool found = false;
            for (auto &prop : props)
            {
                if (strcmp(name, prop.extensionName) == 0)
                {
                    found = true;
                    break;
                }
            }
            if (!found)
            {
                throw std::runtime_error("desired vkDevice-extension not found: " + std::string(name));
            }
        }
    }

}